import os
import numpy as np
from easy_io import read_mat_file, write_pkl_file
import global_config


candidate_arr = read_mat_file('/data_4t/Kaggle/candidates/lidc_candis.mat')
candidates = []
for s in candidate_arr:
    scanid = str(s[0][0])
    assert s[1].ndim == 2
    for c in s[1]:
        c = np.asarray(c[:4], 'int')
        candidates.append(dict(scanid=scanid, center=c[:3], diameter=int(c[3])))

for i in range(len(candidates)):
    candidates[i].update(index=i, path=str(i))

fold_list = np.load('/data_4t/Kaggle/fold_list.npy')
divisions = {}
for i in range(len(fold_list)):
    scans = fold_list[i]
    divisions.update((s, i) for s in scans)
for c in candidates:
    c.update(fold=divisions[c['scanid']])

label_dict = np.load('/data_4t/Kaggle/lidc&kaggle/label_dict3.npy').item()
noduleids = {}
centers = {}
radius = {}
maligns = {}
sources = {}
noduleid = 0
for scanid in sorted(label_dict.keys()):
    nodules = label_dict[scanid]
    if len(nodules) > 0:
        noduleids[scanid] = np.arange(noduleid, noduleid + len(nodules))
        noduleid += len(nodules)
        bbox = np.asarray([n['bbox'] for n in nodules], 'float')
        assert bbox.ndim == 2 and bbox.shape[-1] == 6
        centers[scanid] = (bbox[:, :3] + bbox[:, 3:] - 1) / 2
        radius[scanid] = (bbox[:, 3:] - bbox[:, :3])
        malign = []
        source = []
        for n in nodules:
            if n.get('attrs', None) is not None:
                malignlst = sorted((a['malignancy'] for a in n['attrs']))
                malign.append(malignlst[len(malignlst) // 2])
                source.append('LIDC')
            else:
                malign.append(n['malign'])
                source.append('kaggle')
        maligns[scanid] = np.asarray(malign, 'float')
        assert all(s == source[0] for s in source)
        sources[scanid] = source[0]

for c in candidates:
    scanid = c['scanid']
    center = np.asarray(c['center'], 'float')
    if scanid in noduleids:
        diff = np.sum(((center - centers[scanid]) / radius[scanid]) ** 2, axis=-1) <= 1.0
        # assert 0 <= np.sum(diff) <= 1, c
        if np.sum(diff) > 1:
            print('Warning! More than one nodule has been found by', c)
            c.update(polarity=1, noduleid=int(noduleids[scanid][diff][0]),
                     malign=float(maligns[scanid][diff][0]), source=sources[scanid])
        elif np.sum(diff) == 1:
            c.update(polarity=1, noduleid=int(noduleids[scanid][diff]),
                     malign=float(maligns[scanid][diff]), source=sources[scanid])
        else:
            c.update(polarity=0, noduleid=-1,
                     malign=-1, source=sources[scanid])
    else:
        source = 'LIDC' if scanid.startswith('LIDC') else 'kaggle'
        c.update(polarity=0, noduleid=-1,
                 malign=-1, source=source)

write_pkl_file(os.path.join(global_config.data_folder, 'lidc_candidates.pkl'), candidates)
